"""
配置文件
"""

from pathlib import Path


class Config:
    """配置类"""

    # 路径设置
    ROOT_DIR = Path(__file__).resolve().parent
    DATA_ROOT = ROOT_DIR / "datasets"
    MODEL_PATH = ROOT_DIR / "models" / "pytorch_model.bin"
    OUTPUT_DIR = ROOT_DIR / "outputs"
    CHECKPOINT_DIR = OUTPUT_DIR / "checkpoints"

    # 模型设置 - 使用ViT-L/16
    MODEL_NAME = "vit_large_patch16_224"  # 更改为ViT-L/16
    INPUT_SIZE = 224  # 输入尺寸
    NUM_FEATURES = 1024  # ViT-L的特征维度
    PRETRAINED_MEAN = [0.5, 0.5, 0.5]
    PRETRAINED_STD = [0.5, 0.5, 0.5]

    # 数据集设置
    NUM_KNOWN_CLASSES = 50  # 第一阶段的类别数
    NUM_NEW_CLASSES = 50  # 第二阶段新增的类别数
    NUM_TOTAL_CLASSES = NUM_KNOWN_CLASSES + NUM_NEW_CLASSES 
    NEW_CLASS_SAMPLES = 30  # 每个新类的样本数（Few-shot）

    # 训练设置 - 基于论文优化
    BATCH_SIZE = 64  # ViT-L需要更小的batch size
    EPOCHS = 5  # 第一阶段训练轮数（避免过拟合）
    INCREMENTAL_EPOCHS = 15  # 第二阶段增量学习轮数（增加到15轮）
    LEARNING_RATE = 0.03  # 第一阶段学习率
    INCREMENTAL_LR = 0.01  # 第二阶段学习率（提高10倍）
    MOMENTUM = 0.9
    WEIGHT_DECAY = 0.0  # 论文中没有使用weight decay
    NUM_WORKERS = 4
    WARMUP_STEPS = 5  # 只需要5步warmup
    GRAD_ACCUMULATION_STEPS = 4  # 减少梯度累积以使更新更频繁
    GRAD_NORM_CLIP = 1.0  # 梯度裁剪

    # 增量学习设置
    USE_KNOWLEDGE_DISTILLATION = True  # 使用知识蒸馏防止遗忘
    DISTILLATION_WEIGHT = 0.4  # 蒸馏损失权重（提高以更好保持旧类）
    CLASSIFICATION_WEIGHT = 0.6  # 分类损失权重
    TEMPERATURE = 4.0  # 蒸馏温度
    REPLAY_BUFFER_SIZE = 2000  # 旧类回放缓冲区大小（增加到1500）

    # 模型微调设置
    FREEZE_LAYERS = 0  # 全模型微调，不冻结任何层
    USE_DATA_AUGMENTATION = False  # 短期训练不使用数据增强

    # 评估设置
    EVAL_INTERVAL = 1
    EARLY_STOPPING_PATIENCE = 2  # 早停耐心（轮数）

    # Mahalanobis设置（保留但不使用）
    MAHALANOBIS_LAYER = "features"  # 使用最后一层特征
    USE_RELATIVE_MAHALANOBIS = True  # 使用Relative Mahalanobis Distance
    NORMALIZE_FEATURES = True  # 特征归一化

    def __init__(self):
        """创建必要的目录"""
        self.OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
        self.CHECKPOINT_DIR.mkdir(parents=True, exist_ok=True)


cfg = Config()